import cv2
from PyQt5 import QtCore, QtGui
import json
import os
import time
import numpy as np
from collections import defaultdict
from datetime import datetime


class Worker(QtCore.QThread):
    image_loaded = QtCore.pyqtSignal(QtGui.QPixmap)
    cur_info_loaded = QtCore.pyqtSignal(str)
    image_info_loaded = QtCore.pyqtSignal(str)

    def __init__(self):
        super().__init__()
        self.parent = None
        self.param = None
        self.model_list = []
        self.has_include = False
        self.has_request = False
        self.all_required_list = []
        self.all_include_list = []
        self.all_exclude_list = []
        self.run_flag = True
        self.pause_time = 0

    def run(self):
        print("run开始:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
        self.__extract()

    def __transferModelClass(self):
        self.model_list = []
        model_list = self.param["model_list"]
        tmp_model_list = []
        self.all_required_list = []
        self.all_include_list = []
        self.all_exclude_list = []
        for model_info in model_list:
            # 开始从新计算类配置
            class_list = model_info["class_list"]
            tmp_mode_info = {"file_name": model_info["file_name"], "file_path": model_info["file_path"],
                             "uuid": model_info["uuid"], "model": model_info["model"]}
            if model_info["used"] == 0:
                tmp_model_list.append(tmp_mode_info)
                continue
            include_list = []
            exclude_list = []
            required_list = []
            for class_info in class_list:
                used = class_info["used"]
                sub_key = "%s_%s" % (model_info["uuid"], class_info["class_name"])
                if used == 1:
                    required_list.append(sub_key)
                elif used == 2:
                    include_list.append(sub_key)
                elif used == 3:
                    exclude_list.append(sub_key)
            if len(include_list) > 0:
                self.all_include_list.extend(include_list)
            if len(exclude_list) > 0:
                self.all_exclude_list.extend(exclude_list)
            if len(required_list) > 0:
                self.all_required_list.extend(required_list)
            tmp_model_list.append(tmp_mode_info)
        self.model_list = tmp_model_list
        self.has_include = len(self.all_include_list) > 0
        self.has_request = len(self.all_required_list) > 0

    """
        目标检测
    """

    def __detection_img(self, cap, start_frame, end_frame, step, lable_dir_path, frame_height, frame_width):
        valid = 0
        invalid = 0
        frame_json = {}
        sub_json_list = []
        for idx in range(start_frame, end_frame, step):
            if not self.run_flag:
                break
            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
            ret, frame = cap.read()
            if ret:
                frame_start_time = int(time.time() * 1000)
                # 开始识别
                identify_img, flag, sub_json_list = self.__identify_img(frame)
                file_name = "frame_%d" % idx
                if flag > -1:
                    valid = valid + 1
                    valid_path = os.path.join(lable_dir_path, "valid")
                    if not os.path.exists(valid_path):
                        os.makedirs(valid_path)
                    cv2.imwrite(os.path.join(valid_path, file_name + ".jpg"), frame)
                    if len(sub_json_list) > 0:
                        frame_json = {"version": "3.0.3", "flags": {}, "imageData": None,
                                      "imagePath": file_name + ".jpg", "imageHeight": frame_height,
                                      "imageWidth": frame_width, "shapes": sub_json_list}
                        with open(os.path.join(valid_path, file_name + ".json"), 'w', encoding='utf-8') as f:
                            json.dump(frame_json, f, indent=2, ensure_ascii=False)
                else:
                    invalid = invalid + 1
                    invalid_path = os.path.join(lable_dir_path, "invalid")
                    if not os.path.exists(invalid_path):
                        os.makedirs(invalid_path)
                    cv2.imwrite(os.path.join(invalid_path, file_name + ".jpg"), frame)
                    # 填空白图
                    identify_img = np.zeros((frame_height,frame_width,3 ), np.uint8)
                self.cur_info_loaded.emit(
                    "  开始帧:{},结束帧:{}\n  步长:{}\n  当前帧:{}({})\n  共计:{}帧有效，{}帧无效".format(start_frame,
                                                                                                         end_frame,
                                                                                                         step, idx, (
                                                                                                             "有效" if flag else "无效"),
                                                                                                         valid,
                                                                                                         invalid))
                # 成功读取帧
                image_data = self.parent.cvMatToQImage(identify_img)
                pixmap = QtGui.QPixmap.fromImage(image_data)
                self.image_loaded.emit(pixmap)
                if flag == -1:
                    continue
                frame_end_time = int(time.time() * 1000)
                diff_time = frame_end_time - frame_start_time
                if diff_time > self.pause_time:
                    continue
                time.sleep((self.pause_time - diff_time) / 1000)

    def __track_model(self,mode_info,frame,track_history):
        results = mode_info["model"].track(frame, persist=True)
        if not results or len(results) == 0 or not results[0].boxes:
            return [], track_history
        boxes = results[0].boxes
        detect_list = []
        for box in boxes:
            x = float(box.xyxy[0][0])
            y = float(box.xyxy[0][1])
            w = float(box.xyxy[0][2])
            h = float(box.xyxy[0][3])
            label = f"{results[0].names[int(box.cls[0])]}"
            if box.id:
                track_id = box.id.int().cpu().tolist()
                track = track_history[track_id[0]]
                track.append((float(x), float(y)))
                if len(track) > 30:
                    track.pop(0)
            detect_list.append(
                {"key": "%s_%s" % (mode_info["uuid"], label), "label": label, "top_x": x, "top_y": y,
                 "foot_x": w,
                 "foot_y": h})
        return detect_list ,track_history

    def __track(self, frame, track_dict):
        all_detect_list = []
        for mode_info in self.model_list:
            uuid = mode_info["uuid"]
            if uuid not in track_dict.keys():
                track_dict[uuid] = defaultdict(lambda: [])
            track_history = track_dict[uuid]
            sub_detect_list,track_history = self.__track_model(mode_info,frame,track_history)
            track_dict[uuid] = track_history
            if len(sub_detect_list) > 0:
                all_detect_list.extend(sub_detect_list)
        tmp_frame = frame.copy()
        all_key_list = [item["key"] for item in all_detect_list]
        result_flg = 1
        if self.has_request and not (self.all_required_list <= all_key_list):
            # 通过检查必选标签列表是否是图像所有标签的子集进行判断，
            # 如果 self.all_required_list 不是 all_key_list，则标识图像不包含所有必选标签
            result_flg = -1
        total_labels = 0
        all_labels = []
        sub_json_list = []
        for detect_info in all_detect_list:
            sub_key = detect_info["key"]
            label, top_x, top_y, foot_x, foot_y = detect_info["label"], detect_info["top_x"], detect_info["top_y"], \
                                                  detect_info["foot_x"], detect_info["foot_y"]
            tag_label = label
            if sub_key in self.all_required_list:
                tag_label = "<font color='green'>%s</font>" % label
            elif sub_key in self.all_include_list:
                tag_label = "<font color='blue'>%s</font>" % label
            elif sub_key in self.all_exclude_list:
                tag_label = "<font color='red'>%s</font>" % label
            all_labels.append(tag_label)

            if sub_key in self.all_exclude_list:
                # 如果有标签在排除列表中，则返回原始图像
                result_flg = -1
            if sub_key in self.all_required_list or sub_key in self.all_include_list or (
                    not self.has_request and not self.has_include):
                total_labels = total_labels + 1
                # 画框,如果标签在必须或包含列表中,则在图像上标记出该标签，如果必填列表或者包含列表为空时，所有标签都标记到图像上
                cv2.rectangle(tmp_frame, (int(top_x), int(top_y)),
                              (int(foot_x), int(foot_y)), (255, 0, 0), 2)
                cv2.putText(tmp_frame, label,
                            (int(top_x) - 10, int(top_y) - 10),
                            cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), 1)
                sub_json_dict = {"label": label, "score": None, "group_id": None, "description": "", "difficult": False,
                                 "shape_type": "rectangle", "flags": {}, "attributes": {}, "kie_linking": []}
                sub_json_point = [[top_x, top_y], [foot_x, top_y], [foot_x, foot_y], [top_x, foot_y]]
                sub_json_dict["points"] = sub_json_point
                sub_json_list.append(sub_json_dict)
        if (self.has_request or self.has_include) and total_labels == 0:
            # 没有选择必选和包含标签的，直接识别后返回
            result_flg = -1
        if result_flg == -1:
            return frame, result_flg, [], track_dict, all_labels
        else:
            return tmp_frame, result_flg, sub_json_list, track_dict, all_labels

    """
        目标追踪
    """

    def __track_img(self, cap, start_frame, end_frame, step, lable_dir_path, frame_height, frame_width):
        ret = cap.isOpened()
        track_dict = {}
        valid = 0
        invalid = 0
        cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
        while ret:
            idx = cap.get(cv2.CAP_PROP_POS_FRAMES)
            if idx > end_frame or not self.run_flag:
                break
            if idx % 10 != 0 and (idx - start_frame) % step != 0:
                ret = cap.grab()
                continue
            print("__track_img开始读取idx:", idx, datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
            ret, frame = cap.read()
            # print("__track_img开始读取idx:", idx, datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
            frame_json = {}
            sub_json_list = []
            if ret:
                # print("__track_img开始识别idx:",idx, datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
                identify_img, flag, sub_json_list, track_dict, all_labels = self.__track(frame, track_dict)
                if (idx - start_frame) % step != 0:
                    continue
                file_name = "frame_%d" % idx
                if flag > -1:
                    self.image_info_loaded.emit(" ".join(set(all_labels)))
                    valid = valid + 1
                    valid_path = os.path.join(lable_dir_path, "valid")
                    if not os.path.exists(valid_path):
                        os.makedirs(valid_path)
                    cv2.imwrite(os.path.join(valid_path, file_name + ".jpg"), frame)
                    if len(sub_json_list) > 0:
                        frame_json = {"version": "3.0.3", "flags": {}, "imageData": None,
                                      "imagePath": file_name + ".jpg", "imageHeight": frame_height,
                                      "imageWidth": frame_width, "shapes": sub_json_list}
                        with open(os.path.join(valid_path, file_name + ".json"), 'w', encoding='utf-8') as f:
                            json.dump(frame_json, f, indent=2, ensure_ascii=False)
                else:
                    invalid = invalid + 1
                    invalid_path = os.path.join(lable_dir_path, "invalid")
                    if not os.path.exists(invalid_path):
                        os.makedirs(invalid_path)
                    cv2.imwrite(os.path.join(invalid_path, file_name + ".jpg"), frame)
                    # 填空白图
                #    identify_img = np.zeros((frame_height, frame_width,3), np.uint8)
                self.cur_info_loaded.emit(
                    "  开始帧:{},结束帧:{}\n  步长:{}\n  当前帧:{}({})\n  共计:{}帧有效，{}帧无效".format(start_frame,
                                                                                                         end_frame,
                                                                                                         step, idx, (
                                                                                                             "有效" if flag else "无效"),
                                                                                                         valid,
                                                                                                         invalid))
                # 成功读取帧
                image_data = self.parent.cvMatToQImage(identify_img)
                pixmap = QtGui.QPixmap.fromImage(image_data)
                self.image_loaded.emit(pixmap)
                # print("__track_img完成识别idx:", idx, datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))

    def __extract(self):
        # 类  型 handle_type "按秒提取":"1","按帧提取":"2","总数提取":"3","全部":"4"
        # 间隔时间 interval_num
        # 开始时间 start_time
        # 结束时间 end_time
        # 模型参数 {"model_path":model_path,"class_list":[class_name]}
        self.__transferModelClass()
        handle_type = int(self.param["handle_type"])
        interval_num = float(self.param["interval_num"])
        self.pause_time = int(float(self.param["pause_time"]) * 1000)
        frame_count = self.parent.frame_count
        check_type = int(self.param["check_type"])
        fps = self.parent.fps
        if handle_type == 1:
            start_time = float(self.param["start_time"])
            start_time = 0 if start_time == '' or start_time == 0 else float(start_time)
            ent_time = float(self.param["end_time"])
            ent_time = 0 if ent_time == '' or ent_time == 0 else float(ent_time)
            start_frame = int(start_time * fps)
            end_frame = frame_count if ent_time == 0 else int(ent_time * fps)
        else:
            start_frame = int(float(self.param["start_time"]))
            end_frame = int(float(self.param["end_time"]))
        step = 1
        if handle_type == 1:
            step = int(fps * interval_num)
        elif handle_type == 2:
            step = int(interval_num)
        elif handle_type == 3:
            step = int(frame_count // interval_num)
        video_path = self.param["video_path"]
        video_dir_path = os.path.dirname(video_path)
        video_file_name = os.path.basename(video_path)
        cap = cv2.VideoCapture(video_path)
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        print("识别开始:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
        if check_type == 1:
            lable_dir_path = os.path.join(video_dir_path, video_file_name.split(".")[0],"detect")
            if not os.path.exists(lable_dir_path):
                os.makedirs(lable_dir_path)
            self.__detection_img(cap, start_frame, end_frame, step, lable_dir_path, frame_height, frame_width)
        else:
            lable_dir_path = os.path.join(video_dir_path, video_file_name.split(".")[0],"track")
            if not os.path.exists(lable_dir_path):
                os.makedirs(lable_dir_path)
            self.__track_img(cap, start_frame, end_frame, step, lable_dir_path, frame_height, frame_width)
        self.exit()

    def __identify_img(self, frame):
        all_detect_list = []
        for mode_info in self.model_list:
            sub_lable_list = self.__predict_and_detect(frame, mode_info)
            if len(sub_lable_list) > 0:
                all_detect_list.extend(sub_lable_list)
        # 开始匹配
        tmp_frame = frame.copy()
        # 先检查必选标签
        all_key_list = [item["key"] for item in all_detect_list]
        result_flg = 1
        if self.has_request and not (self.all_required_list <= all_key_list):
            # 通过检查必选标签列表是否是图像所有标签的子集进行判断，
            # 如果 self.all_required_list 不是 all_key_list，则标识图像不包含所有必选标签
            result_flg = -1
        total_labels = 0
        all_labels = []
        sub_json_list = []
        for detect_info in all_detect_list:
            sub_key = detect_info["key"]
            label, top_x, top_y, foot_x, foot_y = detect_info["label"], detect_info["top_x"], detect_info["top_y"], \
                                                  detect_info["foot_x"], detect_info["foot_y"]
            tag_label = label
            if sub_key in self.all_required_list:
                tag_label = "<font color='green'>%s</font>" % label
            elif sub_key in self.all_include_list:
                tag_label = "<font color='blue'>%s</font>" % label
            elif sub_key in self.all_exclude_list:
                tag_label = "<font color='red'>%s</font>" % label
            all_labels.append(tag_label)

            if sub_key in self.all_exclude_list:
                # 如果有标签在排除列表中，则返回原始图像
                result_flg = -1
            if sub_key in self.all_required_list or sub_key in self.all_include_list or (
                    not self.has_request and not self.has_include):
                total_labels = total_labels + 1
                # 画框,如果标签在必须或包含列表中,则在图像上标记出该标签，如果必填列表或者包含列表为空时，所有标签都标记到图像上
                cv2.rectangle(tmp_frame, (int(top_x), int(top_y)),
                              (int(foot_x), int(foot_y)), (255, 0, 0), 2)
                cv2.putText(tmp_frame, label,
                            (int(top_x) - 10, int(top_y) - 10),
                            cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), 1)
                sub_json_dict = {"label": label, "score": None, "group_id": None, "description": "", "difficult": False,
                                 "shape_type": "rectangle", "flags": {}, "attributes": {}, "kie_linking": []}
                sub_json_point = [[top_x, top_y], [foot_x, top_y], [foot_x, foot_y], [top_x, foot_y]]
                sub_json_dict["points"] = sub_json_point
                sub_json_list.append(sub_json_dict)
        if (self.has_request or self.has_include) and total_labels == 0:
            # 没有选择必选和包含标签的，直接识别后返回
            result_flg = -1
        self.image_info_loaded.emit(" ".join(set(all_labels)))
        if result_flg == -1:
            return frame, result_flg, []
        else:
            return tmp_frame, result_flg, sub_json_list

    """循环识别图像，提取图像中所有的标签"""

    def __predict_and_detect(self, img, mode_info):
        results = mode_info["model"].predict(img, conf=0.5)
        sub_list = []
        for result in results:
            for box in result.boxes:
                label = f"{result.names[int(box.cls[0])]}"
                top_x = float(box.xyxy[0][0])
                top_y = float(box.xyxy[0][1])
                foot_x = float(box.xyxy[0][2])
                foot_y = float(box.xyxy[0][3])
                sub_list.append(
                    {"key": "%s_%s" % (mode_info["uuid"], label), "label": label, "top_x": top_x, "top_y": top_y,
                     "foot_x": foot_x,
                     "foot_y": foot_y})
        return sub_list

    def load_data(self, parent, param):
        print("load_data开始:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
        self.parent = parent
        self.param = param
        self.run_flag = True
        self.start()

    def stop(self):
        self.run_flag = False
        print("程序即将停止")
